import numpy as np
from joblib import load

class Classification:
    def __init__(self):
        # 加载训练好的模型、标准化器和标签编码器（指定正确路径：./saved_models/）
        self.knn = load('./saved_models/knn_model.joblib')
        self.scaler = load('./saved_models/scaler.joblib')
        self.le = load('./saved_models/label_encoder.joblib')

    def get_predicter_category(self, x):
        """
        获取预测的分类
        :param x: 包含特征的字典数据
        :return: 预测的分类标签
        """
        # 确保特征顺序与训练时一致
        features_order = [
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ]
        # 提取特征并转换为数组
        new_value = np.array([x[feature] for feature in features_order])
        # 标准化新数据
        new_scaled = self.scaler.transform(new_value.reshape(1, -1))  # 调整为二维数组以适配scaler
        # 预测分类标签（数字编码）
        predicted_label = self.knn.predict(new_scaled)
        # 解码为文本分类
        predicted_category = self.le.inverse_transform(predicted_label)[0]
        return predicted_category

if __name__ == '__main__':
    # 初始化分类器
    cl = Classification()
    # 准备新数据
    new_data = {
        'eps': 1.59,
        'total_revenue_ps': 19.4516,
        'undist_profit_ps': 7.0876,
        'gross_margin': 4063940000,
        'fcff': -279758000,
        'fcfe': -456111000,
        'tangible_asset': 9468120000,
        'bps': 11.5354,
        'grossprofit_margin': 23.4565,
        'npta': 8.3336
    }
    # 预测分类
    predicted_category = cl.get_predicter_category(new_data)
    print(f"预测的分类: {predicted_category}")